#ifndef TEST_Cifar10_MPI_BGV_CIFAR10
#define TEST_Cifar10_MPI_BGV_CIFAR10

#include "Cifar10/mpi_Cifar10_base.hpp"
#include "crypto_helper.hpp"
#include "keys.h"

#define T_NUM 4
#define T_BATCHSIZE 128
#define conv_size 25
#define FC_IN 3
#define HCONV_IN 100
#define m_nblk 1
#define n_nblk 2
#define k_nblk 2

// #define FC_IN 3
#define mp 1
#define np 2
#define kp 32

// #define k_nblk 2
// #define conv_size 121

enum AXIS { M, N, K, AXIS_NUM };

static constexpr double SCALE = 1024.0;

static constexpr unsigned long MATRIX_DIM_M = 128;
static constexpr unsigned long MATRIX_DIM_N = 128;
static constexpr unsigned long MATRIX_DIM_K = 128;
static constexpr unsigned long HYPERCUBE_DIM_X =
    MAX(MATRIX_DIM_M, MATRIX_DIM_K);
static constexpr unsigned long HYPERCUBE_DIM_Y =
    MAX(MATRIX_DIM_K, MATRIX_DIM_N);
// static constexpr unsigned long hmm_grid_size_[AXIS_NUM] = {1, 2, 1};

// std::vector<helib::PubKey* > global_pk;

// double bcast_pubkey(int root, int rank, MPI_Comm comm) {
//   std::string pk_str;
//   std::stringstream pk_ss;
//   auto bgv_client = static_cast<mpi_bgv_hCifar10_client *>(
//       const_cast<mpi_h_Cifar10_client<helib::BGV> *>(client_));

//   if (rank == root) {
//     global_pk->writeTo(pk_ss);
//     pk_str = pk_ss.str();
//   }
//   MPI_Bcast_String(pk_str, root, comm);
//   if (rank != root) {
//     pk_ss << pk_str;
//   }
//   return pk_str.size() / 1024.0 / 1024.0;
// }

static cube_signature MNIST_HYPERCUBE({HYPERCUBE_DIM_X, HYPERCUBE_DIM_Y});

class mpi_bgv_hCifar10_client : public mpi_h_Cifar10_client<helib::BGV> {
private:
  mpi_bgv_hCifar10_client() = delete;

  helib::Ctxt init_hx(const torch::TensorAccessor<float, 2> &a,
                      int t_index) const;

  helib::Ctxt init_hx_single(const torch::TensorAccessor<float, 2> &a,
                             int t_index) const;

  helib::Ctxt init_hconv(const torch::TensorAccessor<float, 3> &a, int asz0,
                         int asz1, int l, int t_index) const;

  helib::Ctxt init_hfc1(const torch::TensorAccessor<float, 3> &a, int asz0,
                        int t_index) const;

  helib::Ctxt init_hfc2(const torch::TensorAccessor<float, 2> &a,
                        int t_index) const;

public:
  mpi_bgv_hCifar10_client(const params<helib::BGV> &params, int crt_num);
  ~mpi_bgv_hCifar10_client() {
    for (auto hmmcc : hmmccs_) {
      delete hmmcc;
    }
    for (auto hmme : hmmes_) {
      delete hmme;
    }
  }

  // double bcast_pubkey(int root, int rank, MPI_Comm comm, helib::SecKey *sk);

  int get_crt_num() const { return crt_num_; }

  const std::vector<helib::Context *> &get_context() const { return contexts_; }

  const std::vector<hypercube_hmmcc *> &get_hmmcc() const { return hmmccs_; }

  const std::vector<hypercube_hmme *> &get_hmme() const { return hmmes_; }
  // Generate encrypted model
  void init_hcnn(encrypted_model<helib::BGV> &ctxt_model,
                 const torch::nn::Conv2d &conv, const torch::nn::Linear &fc1,
                 const torch::nn::Linear &fc2) const override;

  // Generate encrypted image
  void init_input(encrypted_input<helib::BGV> &ctxt_input,
                  const torch::Tensor &filter,
                  const torch::Tensor &x) const override;

  void init_input_single(std::vector<std::string> &send, const torch::Tensor &a,
                         int dst) const override;
  void init_input_single2(std::vector<std::string> &send,
                          const torch::Tensor &a, int dst) const override;
  void recover_result(at::Tensor &out,
                      const std::vector<std::vector<std::string>> &in,
                      MPI_Comm comm, int root) const override;

private:
  ENC key = ENC::PUBKEY;

  int crt_num_;
  NTL::ZZ T_;
  std::vector<NTL::ZZ> ti_;

  std::vector<hypercube_hmmcc *> hmmccs_;
  std::vector<helib::Context *> contexts_;
  std::vector<hypercube_hmme *> hmmes_;
};

class mpi_bgv_hCifar10_server : public mpi_h_Cifar10_server<helib::BGV> {
private:
  mpi_bgv_hCifar10_server() = delete;

public:
  std::shared_ptr<helib::Context>
  build_context(const params<helib::BGV> &params, long m);

  mpi_bgv_hCifar10_server(const hmm_status<shmm_engine> &status, MPI_Comm comm,
                          int crt_num, int conv_size0,
                          const params<helib::BGV> &params);

  mpi_bgv_hCifar10_server(const hmm_status<shmm_engine> &status, MPI_Comm comm,
                          int crt_num, int conv_size0,
                          const params<helib::BGV> &params,
                          const std::vector<hypercube_hmmcc *> &hmmcc);
  ~mpi_bgv_hCifar10_server() {
    for (auto hmme : hmmes_) {
      delete hmme;
    }
    for (auto pk : pk_) {
      delete pk;
    }
  }

  const std::vector<helib::Context *> &get_context() const { return contexts_; }

  int get_crt_num() const { return crt_num_; }

  const std::vector<hypercube_hmme *> &get_hmme() const { return hmmes_; }
  // const std::vector<helib::PubKey *> &get_pk() const { return pk_; }

  const std::vector<hypercube_hmmcs *> &get_hmmcs() const { return hmmcss_; }

  const helib::PubKey &get_pk(int i) const { return *pk_[i]; }

  void forward_test(
      std::vector<std::vector<std::string>> &out,
      std::vector<
          std::vector<std::vector<std::vector<std::vector<std::string>>>>>
          &buf_input,
      const std::vector<std::vector<std::vector<std::vector<helib::Ctxt>>>> &hconv,
      const std::vector<std::vector<helib::Ctxt>> &hfc1,
      const std::vector<helib::Ctxt> &hfc2) const override;

  void report_metrics(std::ostream &ss = std::cout) const override;

  void send_input_p2p(
      const encrypted_input<helib::BGV> &ctxt_in, MPI_Comm comm,
      std::vector<
          std::vector<std::vector<std::vector<std::vector<std::string>>>>> &buf,
      int root) const override;
  // void
  // send_input_p2p_single(const std::vector<std::string> &send, MPI_Comm comm,
  //                       std::vector<std::vector<std::vector<std::string>>>
  //                       &buf, int root, int dst, std::vector<int>
  //                       &buffer_attached_size, std::vector<char *>
  //                       &buffer_attached, int sendbatchsize, MPI_Request
  //                       *reqs) const override;
  void send_input_p2p_single(
      std::vector<std::string> &send, MPI_Comm comm,
      std::vector<
          std::vector<std::vector<std::vector<std::vector<std::string>>>>> &buf,
      int root, int dst) const override;
  void send_hconv_p2p(
      const std::vector<std::vector<std::vector<std::vector<helib::Ctxt>>>>
          &ctxt,
      MPI_Comm comm,
      std::vector<std::vector<std::vector<std::vector<std::string>>>> &buf,
      int root, int rank, int crt_num) const override;
  void back_honv_p2p(const std::vector<std::vector<std::string>> &ctxt,
                     MPI_Comm comm,
                     std::vector<std::vector<std::vector<std::string>>> &buf,
                     int root, int rank, int crt_num) const override;
  void send_hfc1_p2p(std::vector<std::vector<helib::Ctxt>> &ctxt, MPI_Comm comm,
                     std::vector<std::string> &buf, int root, int rank,
                     int crt_num) const override;
  void back_hfc1_p2p(const std::vector<std::string> &ctxt, MPI_Comm comm,
                     std::vector<std::vector<std::string>> &buf, int root,
                     int rank, int crt_num) const override;

  void bcast_hfc2_client(const std::vector<helib::Ctxt> &ctxt, MPI_Comm comm,
                         std::vector<std::string> &buf, int root,
                         int rank) const override;

  void bcast_hconv_server(
      const std::vector<std::vector<std::vector<std::vector<std::string>>>> &in,
      MPI_Comm comm,
      std::vector<std::vector<std::vector<std::vector<std::string>>>> &buf,
      int root, int rank, int crt_num) const override;
  void
  back_hconv_server(const std::vector<std::vector<std::string>> &ctxt,
                    MPI_Comm comm,
                    std::vector<std::vector<std::vector<std::string>>> &buf,
                    int root, int rank, int crt_num) const override;

  void bcast_hfc1_server(const std::vector<std::string> &in, MPI_Comm comm,
                         std::vector<std::string> &buf, int root, int rank,
                         int crt_num) const override;
  void
  back_input_server(const std::vector<std::vector<std::string>> &ctxt,
                    MPI_Comm comm,
                    std::vector<std::vector<std::vector<std::string>>> &buf,
                    int root, int rank, int crt_num) const override;

  void hconv_server(
      std::vector<std::vector<std::vector<helib::Ctxt>>> &out,
      const std::vector<
          std::vector<std::vector<std::vector<std::vector<helib::Ctxt>>>>> &in,
      const std::vector<std::vector<std::vector<std::vector<helib::Ctxt>>>> &hconv) const;
  void back_hconv(const std::vector<helib::Ctxt> &out,
                  std::vector<std::vector<helib::Ctxt>> &ctxt) const;
  void
  hsqr1_server(std::vector<std::vector<std::vector<helib::Ctxt>>> &ctxt) const;
  void hfc1_server(std::vector<std::vector<std::vector<helib::Ctxt>>> &ctxt,
                   const std::vector<std::vector<helib::Ctxt>> &hfc1,
                   std::vector<std::vector<helib::Ctxt>> &out, MPI_Comm comm,
                   int rank, int root) const;
  void back_hfc1(const std::vector<helib::Ctxt> &out,
                 std::vector<std::vector<helib::Ctxt>> &ctxt,
                 MPI_Comm comm) const;
  void hsqr2_server(std::vector<std::vector<helib::Ctxt>> &ctxt) const;
  void hfc2_server(std::vector<std::vector<helib::Ctxt>> &ctxt,
                   std::vector<std::vector<std::string>> &buf,
                   const std::vector<helib::Ctxt> &hfc2, MPI_Comm comm,
                   int rank, int root) const;

public:
  std::vector<helib::PubKey *> pk_;

private:
  std::vector<hypercube_hmmcs *> hmmcss_;

  std::vector<metric> metrics_;
  std::vector<helib::Context *> contexts_;
  int crt_num_;
  std::vector<hypercube_hmme *> hmmes_;
};

#endif /* TEST_Cifar10_BGV_MNIST_HELPER */
